package proxy

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/binary"
	"errors"
	"gitee.com/gcom/ths_socket/socks5"
	"log"
	"math/big"
	"strconv"
	"strings"
	"time"
)

var tlsHostCache, _ = NewLRU(500, nil)

func createTLSConfig(host string, cert *tls.Certificate) *tls.Config {
	ix := strings.Index(host, ":")
	if ix != -1 {
		host = host[:ix]
	}

	if v, ok := tlsHostCache.Get(host); ok {
		return v.(*tls.Config)
	}

	signedCert, err := signCert(cert, host)
	if err != nil {
		log.Printf("error signing certificate: %s", err.Error())
		return nil
	}

	cfg := &tls.Config{
		Certificates:       []tls.Certificate{*signedCert},
		ClientSessionCache: tls.NewLRUClientSessionCache(1024),
	}
	tlsHostCache.Add(host, cfg)
	return cfg

}

func signCert(ca *tls.Certificate, host string) (*tls.Certificate, error) {
	x509ca, err := x509.ParseCertificate(ca.Certificate[0])
	if err != nil {
		return nil, err
	}

	now := time.Now()
	notBefore := now.AddDate(0, 0, -1)
	notAfter := now.AddDate(1, 0, 0)
	if err != nil {
		return nil, err
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	template := &x509.Certificate{
		SerialNumber:          serialNumber,
		Issuer:                x509ca.Subject,
		Subject:               pkix.Name{Organization: []string{"Socks5 mitm proxy"}},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}
	template.DNSNames = append(template.DNSNames, host)
	template.Subject.CommonName = host

	pkey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}

	cert, err := x509.CreateCertificate(rand.Reader, template, x509ca, pkey.Public(), ca.PrivateKey)
	if err != nil {
		return nil, err
	}

	return &tls.Certificate{
		Certificate: [][]byte{cert, ca.Certificate[0]},
		PrivateKey:  pkey,
	}, nil
}

func resolveHost(conn socks5.BufferedConn) (string, error) {
	// 5 bytes: tls header
	// 1 bytes: handshake type
	// 3 bytes: handshake msg length
	// 2 bytes: client tls version
	// 32 bytes: random number
	// 1 bytes: session id length
	buf, _ := conn.Peek(9)
	recordLen := int(binary.BigEndian.Uint16(buf[3:5]))
	handType := buf[5]
	if handType != 0x01 {
		return "", errors.New("Not a tls client hello handshake type: " + strconv.Itoa(int(handType)))
	}
	buf, err := conn.Peek(recordLen)
	if err != nil {
		return "", err
	}
	// now at sessionId length byte
	current := 5 + 38
	// add sessionId length byte and sessionId bytes
	current += 1 + int(buf[current])
	// add cipher length bytes and cipher bytes
	current += 2 + int(binary.BigEndian.Uint16(buf[current:current+2]))
	// add compression method byte and compression bytes
	current += 1 + int(buf[current])
	// add extensions length bytes
	current += 2

	// read extensions
	for current < recordLen {
		et := int(binary.BigEndian.Uint16(buf[current : current+2]))
		el := int(binary.BigEndian.Uint16(buf[current+2 : current+4]))
		if et == 0x00 {
			sn := buf[current+4 : current+4+el]
			ix := 0
			for ix < len(sn) {
				l := int(binary.BigEndian.Uint16(sn[ix : ix+2]))
				if sn[ix+2] == 0x00 {
					//l2 := int(binary.BigEndian.Uint16(sn[ix+3 : ix+5]))
					return string(sn[ix+5:]), nil
				}
				ix += l + 2
			}
		}

		current += 4 + el
	}
	return "", errors.New("no hostname found in client hello handshake")
}
